#include "kernel_operator.h"
#include "lib/matmul_intf.h"

using namespace AscendC;

class NLLLossNew {
 public:
  	__aicore__ inline NLLLossNew() {}
  	__aicore__ inline void Init(GM_ADDR x, GM_ADDR target, GM_ADDR weight,
                              GM_ADDR y, GM_ADDR workspace, uint32_t reduction,
                              int32_t ignoreIndex, uint32_t sizeC,
                              uint32_t coreRowNum, uint32_t coreRowMod ) {
	
    this->reduction = reduction;
    this->ignoreIndex = ignoreIndex;
	this->sizeC = sizeC;
	this->rowSmall = coreRowNum;
	this->rowBig = rowSmall + 1;
	this->rowMod = coreRowMod;
	this->coreNum = GetBlockNum();

	this->coreIndex = GetBlockIdx();
	this->beginRow = (coreIndex < rowMod) ? (rowBig * this->coreIndex) : (rowSmall * this->coreIndex + rowMod);
	this->coreRow = (coreIndex < rowMod) ? rowBig : rowSmall;

    weightGm.SetGlobalBuffer((__gm__ DTYPE_WEIGHT*)weight,
                             sizeC * sizeof(DTYPE_WEIGHT));
    this->roundSizeC = ((this->sizeC+7) / 8) * 8;
    this->roundCoreRow = ((this->coreRow+7) / 8) * 8;
    this->roundCoreNum =  ((this->coreNum+7) / 8) * 8;
    pipe.InitBuffer(AllWeightBuf, roundSizeC * sizeof(DTYPE_WEIGHT));

    targetGm.SetGlobalBuffer(
        (__gm__ DTYPE_TARGET*)target + this->beginRow,
        this->coreRow * sizeof(DTYPE_TARGET));
	
	pipe.InitBuffer(AllTargetBuf, roundCoreRow * sizeof(DTYPE_TARGET));
	// pipe.InitBuffer(AllXBuf, roundCoreRow * sizeof(DTYPE_TARGET));

	// pipe.InitBuffer(inQueueTarget, BUFFER_NUM, ITER_N * sizeof(DTYPE_TARGET));

	pipe.InitBuffer(GatherWeightBuf, roundCoreRow * sizeof(DTYPE_WEIGHT));
	// this->gatherWeight = GatherWeightBuf.Get<DTYPE_WEIGHT>();

	pipe.InitBuffer(InitIndexBuf, roundCoreRow * sizeof(DTYPE_TARGET));

	// pipe.InitBuffer(IndexBuf, this->coreRow * sizeof(DTYPE_TARGET));
	xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->beginRow * this->sizeC,
                        this->coreRow * this->sizeC * sizeof(DTYPE_X));
	

	pipe.InitBuffer(GatherXBuf, roundCoreRow * sizeof(DTYPE_X));
	pipe.InitBuffer(CopyInXBuf, roundCoreRow * sizeof(DTYPE_X) * 8);

	pipe.InitBuffer(ReduceXBuf, 32);
	pipe.InitBuffer(ReduceWBuf, 32);
	pipe.InitBuffer(WorkBuf, (roundCoreRow + roundCoreNum) * sizeof(DTYPE_X));
 
	pipe.InitBuffer(FinalXBuf, this->roundCoreNum * sizeof(DTYPE_X));
	pipe.InitBuffer(FinalWBuf, this->roundCoreNum * sizeof(DTYPE_X));

    yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, sizeof(DTYPE_Y));

	workGm.SetGlobalBuffer((__gm__ DTYPE_X*)workspace, this->roundCoreNum * 2 * sizeof(DTYPE_X));
	// workWGm.SetGlobalBuffer(
	// (__gm__ DTYPE_WEIGHT*)workspace + this->roundCoreNum+ 32,
	// this->roundCoreNum * sizeof(DTYPE_WEIGHT));
	}

	__aicore__ inline void process() {
		EachCoreProcess();
		SyncAll();
		if (this->coreIndex != 0) return;
		FinalCoreProcess();
	}

	__aicore__ inline void EachCoreProcess() {
		/*
			copy in weight
		*/
		auto allWeight =AllWeightBuf.Get<DTYPE_WEIGHT>();
		DataCopy(allWeight, weightGm, roundSizeC);
		PipeBarrier<PIPE_ALL>();

		/*
			copy in target
		*/
		auto allTarget = AllTargetBuf.Get<DTYPE_TARGET>();
		DataCopy(allTarget, targetGm, roundCoreRow);
		PipeBarrier<PIPE_ALL>();

		/*
			gather weight
		*/
		auto indexLocal = InitIndexBuf.Get<DTYPE_TARGET>();
		auto gatherWeight = GatherWeightBuf.Get<DTYPE_WEIGHT>();
		Muls(indexLocal, allTarget, 4 , coreRow);
		auto indexLocalUint = indexLocal.ReinterpretCast<uint32_t>();
		PipeBarrier<PIPE_ALL>();
		Gather(gatherWeight, allWeight, indexLocalUint, (uint32_t)0, this->coreRow);

		/*
			gather x
		*/
		auto copyInX = CopyInXBuf.Get<DTYPE_X>();
		auto gatherX = GatherXBuf.Get<DTYPE_X>();
		
        AscendC::DataCopyPadParams padParams{false, 0, 0, 0};
        AscendC::DataCopyParams copyParams{1, 4, 0, 0};

        for (int i = 0; i < coreRow; i++) {
            AscendC::DataCopyPad(copyInX[i*8], xGm[i * sizeC + allTarget.GetValue(i)], copyParams, padParams);
        }
		CreateVecIndex(indexLocal, (int32_t)0, coreRow);
		Muls(indexLocal, indexLocal, 32 , coreRow);
		PipeBarrier<PIPE_ALL>();
		Gather(gatherX, copyInX, indexLocalUint, (uint32_t)0, this->coreRow);

		/*
			Compute
		*/
		auto workLocal = WorkBuf.Get<DTYPE_X>();
		auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
		auto reduceWLocal = ReduceWBuf.Get<DTYPE_WEIGHT>();
		if (this->reduction == 2) {	//mean
			ReduceSum(reduceWLocal, gatherWeight, workLocal, this->coreRow);
		}

		Mul(gatherX, gatherX, gatherWeight, this->coreRow);
		// for (int i = 0; i < this->coreRow; i++) {
		// 	printf("i: %d %f %f\n",i, gatherX.GetValue(i),gatherWeight.GetValue(i));
		// }
		ReduceSum(reduceXLocal, gatherX, workLocal, this->coreRow);
    	printf("reducexLocal: %f \n",reduceXLocal.GetValue(0));

		PipeBarrier<PIPE_ALL>();

		/*
			CopyToWorkSpace
		*/
	    AscendC::DataCopyExtParams copyExtParams{1, 4, 0, 0, 0};
   		AscendC::DataCopyPad(workGm[this->coreIndex], reduceXLocal, copyExtParams);
		if (this->reduction == 2)  // mean
			AscendC::DataCopyPad(workGm[this->coreIndex + this->roundCoreNum], reduceWLocal, copyExtParams);
		// DataCopy(workWGm[this->coreIndex * ITER_N], LocaloutW, ITER_N);
		PipeBarrier<PIPE_ALL>();

	}

	__aicore__ void FinalCoreProcess() {
		auto finalXLocal = FinalXBuf.Get<DTYPE_X>();
		auto finalWLocal = FinalWBuf.Get<DTYPE_WEIGHT>();

		// for (int i = 0; i < 2000; i++)
		// 	finalXLocal.GetValue(i % 8);
		DataCopy(finalXLocal, workGm, this->roundCoreNum);
		if (this->reduction == 2) {	//mean
			DataCopy(finalWLocal, workGm[this->roundCoreNum], this->roundCoreNum);
		}

		PipeBarrier<PIPE_ALL>();
		for (int i = 0; i < this->coreRow; i++)
			printf("final %f %f\n", finalXLocal.GetValue(i),finalWLocal.GetValue(i));

		auto reduceXLocal = ReduceXBuf.Get<DTYPE_X>();
		auto reduceWLocal = ReduceWBuf.Get<DTYPE_WEIGHT>();
		auto workLocal = WorkBuf.Get<DTYPE_X>();
		float p = -1;
		ReduceSum(reduceXLocal, finalXLocal, workLocal, this->coreNum);

		Muls(reduceXLocal,reduceXLocal,p,1);
		if (this->reduction == 2) {	//mean
			ReduceSum(reduceWLocal, finalWLocal, workLocal, this->coreNum);
			Div(reduceXLocal,reduceXLocal,reduceWLocal,1);
		}
	    AscendC::DataCopyExtParams copyExtParams{1, 4, 0, 0, 0};
		AscendC::DataCopyPad(yGm, reduceXLocal, copyExtParams);
		// DataCopy(yGm, workWGm, (this->coreNum + 7)/ 8* 8);

	}


private:
	uint32_t reduction, sizeC;
	int32_t ignoreIndex;
	uint32_t rowSmall, rowBig, rowMod, coreNum,roundSizeC,roundCoreRow,roundCoreNum;
	uint32_t beginRow,coreIndex,coreRow;

	AscendC::TPipe pipe;
	AscendC::GlobalTensor<DTYPE_X> xGm;
	AscendC::GlobalTensor<DTYPE_TARGET> targetGm;
	AscendC::GlobalTensor<DTYPE_WEIGHT> weightGm;
	AscendC::GlobalTensor<DTYPE_Y> yGm, workGm;

	// TQue<QuePosition::VECIN, 2>  
	TBuf<QuePosition::VECCALC>  AllTargetBuf, AllWeightBuf;
	TBuf<QuePosition::VECCALC> InitIndexBuf,IndexBuf,GatherWeightBuf,GatherXBuf, CopyInXBuf;
	TBuf<QuePosition::VECCALC> ReduceXBuf,ReduceWBuf,WorkBuf;
	TBuf<QuePosition::VECCALC> FinalXBuf,FinalWBuf;

	LocalTensor<DTYPE_WEIGHT> allWeight,gatherWeight;
	LocalTensor<DTYPE_TARGET> allTarget;
	LocalTensor<DTYPE_TARGET> initIndex;
	LocalTensor<DTYPE_X> gatherX;
	
};